package autosql

import (
	"bufio"
	"fmt"
	"go/format"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/text/cases"
	"golang.org/x/text/language"
	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/driver/sqlserver"
	"gorm.io/gen"
	"gorm.io/gorm"
)

// MySQL 连接
func ConnMysql(username, password, host, port, dbname string) (*gorm.DB, error) {
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local", username, password, host, port, dbname)
	db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{})
	if err != nil {
		return nil, err
	}
	return db, nil
}

// PostgreSQL 连接
func ConnPostgres(username, password, host, port, dbname string) (*gorm.DB, error) {
	dsn := fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%s sslmode=disable TimeZone=Asia/Shanghai", host, username, password, dbname, port)
	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
	if err != nil {
		return nil, err
	}
	return db, nil
}

// SQLite 连接
func ConnSQLite(filepath string) (*gorm.DB, error) {
	// SQLite 数据库使用文件路径连接
	db, err := gorm.Open(sqlite.Open(filepath), &gorm.Config{})
	if err != nil {
		return nil, err
	}
	return db, nil
}

// SQL Server 连接
func ConnSQLServer(username, password, host, port, dbname string) (*gorm.DB, error) {
	dsn := fmt.Sprintf("sqlserver://%s:%s@%s:%s?database=%s", username, password, host, port, dbname)
	db, err := gorm.Open(sqlserver.Open(dsn), &gorm.Config{})
	if err != nil {
		return nil, err
	}
	return db, nil
}

// 生成数据库中指定表的结构体（在mod同级目录调用）
func GenStruct(db *gorm.DB, tableName string) {
	// 初始化 GORM GEN 并在当前目录生成 model 包
	g := gen.NewGenerator(gen.Config{
		OutPath: "./query", // 注意query是必须的，但不会作为目录
	})
	// 使用已经连接的数据库
	g.UseDB(db)
	// 生成指定表的结构体
	g.GenerateModel(tableName)
	// 执行生成
	g.Execute()
	// 生成外键信息
	foreignKeys, err := genForeignKeys(db)
	if err != nil {
		fmt.Println(err)
		return
	}
	//如果无外键，不做处理
	if len(foreignKeys) == 0 {
		return
	}
	// 修改生成的文件
	err = modifyGenFiles(foreignKeys)
	if err != nil {
		fmt.Println(err)
		return
	}
}

// 生成数据库中所有表的结构体（在mod同级目录调用）
func GenAllStructs(db *gorm.DB) {
	// 初始化 GORM GEN 并在当前目录生成 model 包
	g := gen.NewGenerator(gen.Config{
		OutPath: "./query", // 注意query是必须的，但不会作为目录
	})
	// 使用已经连接的数据库
	g.UseDB(db)
	// 生成所有表的结构体
	g.GenerateAllTable()
	// 执行生成
	g.Execute()
	// 生成外键信息
	foreignKeys, err := genForeignKeys(db)
	if err != nil {
		fmt.Println(err)
		return
	}
	//如果无外键，不做处理
	if len(foreignKeys) == 0 {
		return
	}
	// 修改生成的文件
	err = modifyGenFiles(foreignKeys)
	if err != nil {
		fmt.Println(err)
		return
	}
}

// 获得数据库中的外键信息
func genForeignKeys(db *gorm.DB) (map[string][]string, error) {
	query := `
	SELECT
		TABLE_NAME, COLUMN_NAME, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME
	FROM
		information_schema.KEY_COLUMN_USAGE
	WHERE
		TABLE_SCHEMA = DATABASE() AND REFERENCED_TABLE_NAME IS NOT NULL;
	`

	rows, err := db.Raw(query).Rows()
	if err != nil {
		err = fmt.Errorf("failed to query foreign key information: %v", err)
		return nil, err
	}

	defer rows.Close()

	foreignKeys := make(map[string][]string)
	for rows.Next() {
		var tableName, columnName, referencedTableName, referencedColumnName string
		err := rows.Scan(&tableName, &columnName, &referencedTableName, &referencedColumnName)
		if err != nil {
			err = fmt.Errorf("failed to scan foreign key information: %v", err)
			return nil, err
		}
		// 生成结构体字段
		structField := fmt.Sprintf("%s *%s `gorm:\"foreignKey:%s;references:%s\" json:\"%s\"`",
			toCamelCase(referencedTableName), toCamelCase(referencedTableName),
			toCamelCase(columnName), toCamelCase(referencedColumnName), referencedTableName)
		foreignKeys[tableName] = append(foreignKeys[tableName], structField)
	}

	return foreignKeys, nil
}

// 将下划线分隔的字符串转换为驼峰命名法
func toCamelCase(input string) string {
	words := strings.Split(input, "_")

	// 使用 cases.Title() 替代 strings.Title
	caser := cases.Title(language.Und)
	for i := range words {
		// 将每个单词的首字母大写
		//遇到id则全部大写
		if words[i] == "id" {
			words[i] = "ID"
		} else {
			words[i] = caser.String(words[i])
		}
	}

	// 拼接所有单词为一个驼峰式字符串
	return strings.Join(words, "")
}

// 修改生成的文件，去掉前三行并替换为一行注释，加入外键信息
func modifyGenFiles(foreignKeys map[string][]string) error {
	// 遍历目标目录下的所有 Go 文件
	return filepath.Walk("./model", func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		// 仅处理 Go 文件
		if !info.IsDir() && strings.HasSuffix(info.Name(), ".go") {
			tableName := strings.TrimSuffix(info.Name(), ".gen.go")
			err := modifyFile(path, foreignKeys[tableName])
			if err != nil {
				return fmt.Errorf("failed to modify file %s: %w", path, err)
			}
			// 格式化
			err = formatFile(path)
			if err != nil {
				return fmt.Errorf("failed to format file %s: %w", path, err)
			}
		}
		return nil
	})
}

// 修改文件，修改前三行并加入外键信息
func modifyFile(filePath string, foreignKeys []string) error {
	// 打开文件
	file, err := os.Open(filePath)
	if err != nil {
		return err
	}
	defer file.Close()

	// 读取文件内容
	var lines []string
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		return err
	}

	// 修改前三行
	lines[0] = "// Code generated by gorm.io/gen. Modified by gitee.com/zht639."
	lines[1] = "// DO NOT EDIT! Next generation will overwrite this file."
	lines[2] = "// DO NOT EDIT! Next generation will overwrite this file."

	// 加入外键信息
	for i := range foreignKeys {
		foreignKeys[i] += " // 外键字段"
	}
	for i := range lines {
		if strings.Contains(lines[i], "struct") {
			// 将外键字段插入到 struct 的下一行
			lines = append(lines[:i+1], append(foreignKeys, lines[i+1:]...)...)
			break
		}
	}

	// 将修改后的内容写回文件
	return writeFile(filePath, lines)
}

// 将修改后的内容写回文件
func writeFile(filePath string, lines []string) error {
	file, err := os.OpenFile(filePath, os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}
	defer file.Close()

	writer := bufio.NewWriter(file)
	for _, line := range lines {
		_, err := writer.WriteString(line + "\n")
		if err != nil {
			return err
		}
	}
	return writer.Flush()
}

// formatFile 格式化给定的 Go 文件
func formatFile(filename string) error {
	// 读取文件内容
	source, err := os.ReadFile(filename)
	if err != nil {
		return err
	}

	// 格式化 Go 源代码
	formattedSource, err := format.Source(source)
	if err != nil {
		return err
	}

	// 写入格式化后的内容到文件
	return os.WriteFile(filename, formattedSource, 0644)
}
